from unittest import TestCase

from algo.colouring import (
    check_colouring,
    find_colouration,
    get_edges,
    get_largest_neighbour,
)


class TestColouring(TestCase):
    def test_get_largest_neighbour(self):
        graph = {
            "A": ["AB", "AC"],
            "B": ["AB", "BC", "BD"],
            "C": ["AC", "BC", "CD"],
            "D": ["BD", "CD"],
        }
        expected = "B"
        results = get_largest_neighbour(graph)
        self.assertEqual(results, expected)

    def test_get_edges(self):
        graph = {
            "A": ["AB", "AC"],
            "B": ["AB", "BC", "BD"],
            "C": ["AC", "BC", "CD"],
            "D": ["BD", "CD"],
        }
        expected = sorted(["AB", "AC", "BC", "CD", "BD"])
        results = get_edges(graph)
        self.assertEqual(results, expected)

    def test_get_edges2(self):
        graph = {
            "A": ["AB"],
            "B": ["AB"],
        }
        expected = sorted(["AB"])
        results = get_edges(graph)
        self.assertEqual(results, expected)


class TestCheckColouring(TestCase):
    def test_check_colouring_fail(self):
        graph = {
            "A": ["AB"],
            "B": ["AB", "BC"],
            "C": ["BC", "CD"],
            "D": ["CD", "DE"],
            "E": ["DE", "EF"],
            "F": ["EF", "FG"],
            "G": ["FG", "GH"],
            "H": ["GH"],
        }
        node_colours = {
            "A": 3,
            "B": 3,
            "C": 2,
            "D": 3,
            "E": 1,
            "F": 2,
            "G": 3,
            "H": 1,
        }
        success = check_colouring(graph, node_colours)
        self.assertTrue(success is False)

    def test_check_colouring(self):
        graph = {
            "A": ["AB"],
            "B": ["AB", "BC"],
            "C": ["BC", "CD"],
            "D": ["CD", "DE"],
            "E": ["DE", "EF"],
            "F": ["EF", "FG"],
            "G": ["FG", "GH"],
            "H": ["GH"],
        }
        node_colours = {
            "A": 3,
            "B": 1,
            "C": 2,
            "D": 3,
            "E": 1,
            "F": 2,
            "G": 3,
            "H": 1,
        }
        success = check_colouring(graph, node_colours)
        self.assertTrue(success)


class TestFindColouration(TestCase):
    def test_find_colouration_empty_graph(self):
        graph = {}
        expected = ({}, 1, 0)
        results = find_colouration(graph)
        self.assertEqual(results, expected)

    def test_find_colouration1(self):
        graph = {
            "A": ["AB"],
            "B": ["AB"],
        }
        expected = (
            {
                "A": 2,
                "B": 1,
            },
            2,
            1,
        )
        results = find_colouration(graph)
        self.assertEqual(results, expected)
        success = check_colouring(graph, results[0])
        self.assertTrue(success)

    def test_find_colouration2(self):
        graph = {
            "A": ["AB", "AC"],
            "B": ["AB"],
            "C": ["AC"],
        }
        expected = (
            {
                "A": 3,
                "B": 2,
                "C": 2,
            },
            3,
            2,
        )
        results = find_colouration(graph)
        self.assertEqual(results, expected)
        success = check_colouring(graph, results[0])
        self.assertTrue(success)

    def test_find_colouration3(self):
        graph = {
            "A": ["AB", "AC"],
            "B": ["AB", "BC", "BD"],
            "C": ["AC", "BC", "CD"],
            "D": ["BD", "CD"],
        }
        expected = (
            {
                "A": 4,
                "B": 3,
                "C": 2,
                "D": 4,
            },
            4,
            3,
        )
        results = find_colouration(graph)
        self.assertEqual(results, expected)
        success = check_colouring(graph, results[0])
        self.assertTrue(success)

    def test_find_colouration4(self):
        graph = {
            "A": ["AB", "AC", "AD", "AE"],
            "B": ["AB", "BC"],
            "C": ["AC", "BC"],
            "D": ["AD"],
            "E": ["EF", "EG", "EH", "AE"],
            "F": ["EF", "FG"],
            "G": ["EG", "FG"],
            "H": ["EH"],
        }
        expected_colours = {
            "A": 5,
            "B": 4,
            "C": 3,
            "D": 4,
            "E": 4,
            "F": 5,
            "G": 3,
            "H": 5,
        }
        expected_number = 5
        expected_neighbours = 4
        node_colours, k_number, d_neighbours = find_colouration(graph)
        self.assertEqual(node_colours, expected_colours)
        self.assertEqual(k_number, expected_number)
        self.assertEqual(d_neighbours, expected_neighbours)
        success = check_colouring(graph, node_colours)
        self.assertTrue(success)

    def test_find_colouration5(self):
        graph = {
            "A": ["AB", "AH"],
            "B": ["AB", "BC"],
            "C": ["BC", "CD"],
            "D": ["CD", "DE"],
            "E": ["DE", "EF"],
            "F": ["EF", "FG"],
            "G": ["FG", "GH"],
            "H": ["AH", "GH"],
        }
        expected = (
            {
                "A": 3,
                "B": 2,
                "C": 3,
                "D": 2,
                "E": 3,
                "F": 2,
                "G": 3,
                "H": 2,
            },
            3,
            2,
        )
        results = find_colouration(graph)
        self.assertEqual(results, expected)
        success = check_colouring(graph, results[0])
        self.assertTrue(success)

    def test_find_colouration6(self):
        graph = {
            "A": ["AB"],
            "B": ["AB", "BC"],
            "C": ["BC", "CD"],
            "D": ["CD", "DE"],
            "E": ["DE", "EF"],
            "F": ["EF", "FG"],
            "G": ["FG", "GH"],
            "H": ["GH"],
        }
        expected = (
            {
                "A": 3,
                "B": 2,
                "C": 3,
                "D": 2,
                "E": 3,
                "F": 2,
                "G": 3,
                "H": 2,
            },
            3,
            2,
        )
        results = find_colouration(graph)
        self.assertEqual(results, expected)
        success = check_colouring(graph, results[0])
        self.assertTrue(success)
